import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
from PIL import Image
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import Config
from datalist import Dataset, data_transform
from model import UNet
from utils import Color_print, WaterMark, denormalize

img_size = 512
best_loss = 1000000


def get_image(image):
    output = (denormalize(image.permute((0, 2, 3, 1)).
                          detach().
                          to('cpu').
                          numpy()) * 255).astype('uint8')
    output = output[0]

    output = Image.fromarray(output).convert('RGB')
    return output


class Nosie2Noisetrain(object):
    def __init__(self):
        self.args = Config()
        print(f"-----------{self.args.project_name}-------------")

        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        if use_cuda:
            torch.cuda.manual_seed(self.args.seed)
            torch.cuda.manual_seed_all(self.args.seed)
        else:
            torch.manual_seed(self.args.seed)

        self.device = torch.device('cuda' if use_cuda else 'cpu')

        kwargs = {"num_workers": 0, "pin_memory": True} if use_cuda else {}

        '''
        构造DataLoader
        '''
        self.train_dataset = Dataset(type='train')
        self.train_loader = DataLoader(self.train_dataset, batch_size=self.args.train_batch_size, shuffle=True,
                                       **kwargs)
        self.test_dataset = Dataset(type='test')
        self.test_loader = DataLoader(self.test_dataset, batch_size=self.args.test_batch_size, shuffle=True,
                                      **kwargs)
        '''
        定义选择模型
        '''
        self.model = UNet().to(self.device)

        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))
            cudnn.enabled = True
            cudnn.benchmark = True
        '''
        根据需要加载与训练的模型权重参数
        '''
        if self.args.resume==False and self.args.pretrained_model:
            model_dict = self.model.state_dict()
            checkpoint = torch.load(self.args.pretrained_model)
            pretrained_dict = checkpoint['model_state_dict']
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if
                               np.shape(model_dict[k]) == np.shape(v)}
            model_dict.update(pretrained_dict)
            self.model.load_state_dict(model_dict, strict=True)
            print("Restoring the weight from pretrained-weight file \nFinished to load the weight")



        '''
        构造loss目标函数
        选择优化器
        学习率变化选择
        '''
        self.criterion = nn.L1Loss()
        self.optimizer = optim.Adam(params=self.model.parameters(),
                                    lr=self.args.lr,
                                    betas=[0.9, 0.99],
                                    eps=1e-8)
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, 1e-2, 1e-5)

        try:
            for epoch in range(1, self.args.epochs + 1):
                self.train(epoch)
                self.test(epoch)
            torch.cuda.empty_cache()
            print("finish model training")
        except KeyboardInterrupt:
            self.savekeyboardinterprut()

    def train(self, epoch):
        self.model.train()
        average_loss = []
        pbar = tqdm(self.train_loader, desc=f'Train Epoch {epoch}/{self.args.epochs}')

        for data, target in pbar:
            data, target = data.to(self.device), target.to(self.device)
            self.optimizer.zero_grad()  # 模型参数梯度清零
            output = self.model(data)
            output = nn.Tanh()(output)
            loss = self.criterion(output, target)
            loss.backward()
            average_loss.append(loss.item())
            self.optimizer.step()
            # img=get_image(output)
            # img.save("result/1.jpg")
            pbar.set_description(f'Train Epoch: {epoch}/{self.args.epochs} loss: {np.mean(average_loss)}')
        self.scheduler.step()

    def test(self, epoch):
        self.model.eval()
        average_loss = []
        pbar = tqdm(self.test_loader, desc=f'Test Epoch {epoch}/{self.args.epochs}')
        with torch.no_grad():
            for data, target in pbar:
                data, target = data.to(self.device), target.to(self.device)
                output = self.model(data)
                output = nn.Tanh()(output)
                loss = self.criterion(output, target)
                average_loss.append(loss.item())
                pbar.set_description(f'Test Epoch: {epoch}/{self.args.epochs} loss: {np.mean(average_loss)}')
        self.save(epoch, average_loss)

    def save(self, epoch, average_loss):
        global best_loss
        '''
         if loss is smaller the last epoch, than save the model
         '''
        if self.args.save and best_loss > np.mean(average_loss):
            best_loss = np.mean(average_loss)
            torch.save({
                'epoch': epoch,
                'model_state_dict': self.model.state_dict(),
                'optimizer_state_dict': self.optimizer.state_dict(),
                'loss': round(np.mean(average_loss), 2)
            }, 'weights/temp/best.pth')
            Color_print("\nmodel saved, the performance if better than before")
        '''
        save the test result for each epoch
        '''
        with torch.no_grad():
            data = Image.open("./000005.jpg")
            data = data.resize((img_size, img_size))
            temp1 = data
            data = WaterMark()(data)

            temp2 = data
            data = data_transform(data)

            data = data.unsqueeze(0)
            data = data.to(self.device)
            output = self.model(data)
            output = nn.Tanh()(output)
            result = get_image(output)

            blank = Image.new(mode='RGB', size=(3 * img_size, img_size), color=(255, 255, 255))

            blank.paste(temp1, (0, 0))
            blank.paste(temp2, (img_size, 0))
            blank.paste(result, (img_size * 2, 0))
            blank.save(f'result/{epoch}.jpg')

    def savekeyboardinterprut(self):
        '''
        if user use ctrl + c to stop the process , the parameters of the model will be saved for next
        time to finetune
        :return:
        '''
        torch.save({
            'model_state_dict': self.model.state_dict(),
        }, 'weights/temp/best.pth')
        Color_print("\nout of the train mode and save the model to the ./weights/best.pth")


Nosie2Noisetrain()
